from openai_util import *
from sys_util import *

def print_red(text):
    print("\033[91m" + text + "\033[0m")  # 使用 ANSI Escape Codes 设置文本颜色为红色

def display_options(data):
    print("该程序具有的功能：")
    for index, (key, value) in enumerate(data.items(), start=1):
        print(f"{index}: {value['description']}")

def show_category_content(category_data):
    print(f"System Content: {category_data['system_content']}")
    print("User Assistant Contents:")
    for content in category_data['user_assistant_contents']:
        print(f"  {content[0]}: {content[1]}")

def get_user_choice(min_num = 1, max_num = 2):
    while True:
        choice = get_single_line_input()
        try:
            choice = int(choice)
            if min_num <= int(choice) <= max_num:
                return int(choice)
            else:
                print("Invalid choice. Please enter a valid category number.")
        except ValueError:
            print("Invalid input. Please enter a valid number.")

@text2exit
def get_single_line_input():
    single_line = input()
    return single_line

def get_multi_line_input(finish_text = "f"):
    multi_line = ""
    while True:
        line = get_single_line_input()
        if line.lower() == finish_text:
            break
        multi_line += line + "\n"
    return multi_line

def content2dic(role, content, delimiter=None):
    # remove possible delimiters in the user's message
    if delimiter is not None:
        content = content.replace(delimiter, "")
        return {"role": role, "content": f"{delimiter}{content}{delimiter}"}
    else:
        return {"role": role, "content": content}

def get_user_content(is_multi_line = False):
    user_content = ""
    finish_text = 'f'

    display_exit()
    if is_multi_line:
        print(f"Enter text, type \'{finish_text}\' to finish: ")
        user_content = get_multi_line_input(finish_text=finish_text)
    else :
        print(f"Enter text: ")
        user_content = get_single_line_input()

    return user_content

def get_user_content2dic(delimiter = "####", is_multi_line = False):
    user_content = ""
    finish_text = 'f'

    display_exit()
    if is_multi_line:
        print(f"Enter text, type \'{finish_text}\' to finish: ")
        user_content = get_multi_line_input(finish_text=finish_text)
    else :
        print(f"Enter text: ")
        user_content = get_single_line_input()

    user_input_dic = content2dic(role="user", content=user_content, delimiter=delimiter)
    return user_input_dic


def update_list(lst, new_item, total_len):
    if len(lst) <= total_len:
        lst.append(new_item)
    else:
        lst.pop(total_len)
        lst.append(new_item)

# 使用装饰器
@rate_limited()
def generate_assistant_content(message_dic, user_input_dic, total_len):

    update_list(message_dic, user_input_dic, total_len)
    print(f"\nGenerating...\n")
    assistant_content = generate_answer_stream(message_dic)
    # assistant_content = generate_answer(message_dic)
    
    # assistant_output_dic = {"role": "assistant", "content": assistant_content}
    # update_list(message_dic, assistant_output_dic)

    return assistant_content


def get_user_input(is_multi_line):
    print("=================================================")
    return get_user_content(is_multi_line=is_multi_line)
